
from langchain_community.vectorstores import Chroma
from langchain_openai.embeddings import AzureOpenAIEmbeddings  # 导入嵌入模型
from langchain_community.document_loaders import PyPDFLoader
from tool import get_azure_endpoint, get_api_key, get_api_version
if __name__ == '__main__':
    embedding = AzureOpenAIEmbeddings(
        azure_endpoint=get_azure_endpoint().rstrip('/'),  # 移除尾部斜杠，只保留基础URL
        model="text-embedding-3-small",  # 重命名为 azure_deployment
        api_key=get_api_key(),
        api_version=get_api_version()
    )
    vectordb_chinese = Chroma(
        persist_directory="./docs/chroma/matplotlib",
        embedding_function=embedding,
    )

    texts_chinese = [
        """毒鹅膏菌(Amanita phalloides)具有大型且引人注目的地上(epigeous)子实体(basidiocarp)蘑菇""",
        """一种具有大型子实体的蘑菇是毒鹅膏菌(Amanita phalloides)。某些品种全白。""",
        """A. phalloides，又名死亡帽，是已知所有蘑菇中最有毒的一种。""",
    ]
    smalldb_chinese = Chroma.from_texts(texts_chinese, embedding=embedding)

    question_chinese = "告诉我关于具有大型子实体的全白色蘑菇的信息"

    texts1 = smalldb_chinese.similarity_search(question_chinese, k=2)
    print("texts1:", texts1)

    texts2 = smalldb_chinese.max_marginal_relevance_search(question_chinese, k=2, fetch_k=3)
    print("texts2:", texts2)

